{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b4da38ad",
   "metadata": {},
   "source": [
    "# 1. Libraries & Sample Data\n",
    "The first step is to load our Python Libraries and download the sample data. The dataset represents Apple stock price (1d bars) for the year 2010"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aa05430",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Python Libraries\n",
    "import math\n",
    "import keras\n",
    "import random\n",
    "import datetime\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from collections import deque\n",
    "from tqdm.notebook import tqdm\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "# for dataframe display\n",
    "pd.set_option(\"display.max_rows\", None)\n",
    "def display_df(df):\n",
    "    # Puts the scrollbar next to the DataFrame\n",
    "    display(HTML(\"<div style='height: 200px; overflow: auto; width: fit-content'>\" + df.to_html() + \"</div>\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa831031",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download Sample Data (cleaned, not normalized, without features)\n",
    "data = pd.read_csv('aapl_2010_3m_CLEAN.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60e8526a",
   "metadata": {},
   "source": [
    "# 2. Train / Test Split\n",
    "Now that we have our our cleaned price dataset, we are ready to feed the data into our model. With this in mind, we select Close as our singular training feature, and split the data ito train and test data (80/20 split)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7557879",
   "metadata": {},
   "outputs": [],
   "source": [
    "# split dataset df into train (80%) and test (20%) datasets\n",
    "dataset = data[['Date', 'Close']]\n",
    "\n",
    "training_rows = int(len(dataset.index)*0.8)\n",
    "train_df = dataset.loc[:training_rows].set_index(\"Date\")\n",
    "test_df = dataset.loc[training_rows+1:].set_index(\"Date\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1dca49e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# display train and test dfs (ensure no overlap)\n",
    "display_df(train_df)\n",
    "display_df(test_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92a6ed95",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert train and test dfs to np arrays with dtype=float\n",
    "X_train = train_df.values.astype(float)\n",
    "X_test = test_df.values.astype(float)\n",
    "# print the shape of X_train to remind yourself how many examples and features are in the dataset\n",
    "X_train.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "714c77a8",
   "metadata": {},
   "source": [
    "# 3. Define the Agent\n",
    "Now that our data is ready to use, we can define the Reinforcement Learning Agent."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abd06b6d",
   "metadata": {},
   "source": [
    "### Define the DQN Model\n",
    "The first step in defining our agent is the Deep Q-Network model definition. For this excercise, we are creating a sequential model with three layers. The first two layers have output shape of 32 and 8, respectively, and a RELU activation. The output layer has an output shape of the size of our action space (buy, sell, hold), and a linear activation. Our Loss function is Mean Squared Error, and our optimizer is Adam with a learning rate of 0.001. Use Keras to build this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1e98591",
   "metadata": {},
   "outputs": [],
   "source": [
    "@keras.saving.register_keras_serializable()\n",
    "# Define DQN Model Architecture\n",
    "class DQN(keras.Model):\n",
    "    def __init__(self, state_size, action_size):\n",
    "    \n",
    "        model = keras.models.Sequential()\n",
    "        #Input Layer\n",
    "        model.add(keras.layers.Dense(units=32, input_dim=state_size, activation=\"relu\"))\n",
    "        #Hidden Layer\n",
    "        model.add(keras.layers.Dense(units=8, activation=\"relu\"))\n",
    "        #Output Layer \n",
    "        model.add(keras.layers.Dense(action_size, activation=\"linear\"))\n",
    "        model.compile(loss=\"mse\", optimizer=keras.optimizers.Adam(learning_rate=0.001))\n",
    "\n",
    "        self.model = model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edabcd7d",
   "metadata": {},
   "source": [
    "### Define Agent Class\n",
    "Now that we have defined our underlying DQN Model, we must define out Reinforcement Learning Agent. The agent initialization is provided for you, you must define an act function, and an expereince replay function. As a reminder, the act function defines how our model will act (buy, hold, or sell) given a certain state. The Experience Replay function tackles catastrophic forgetting in our training process, by maintaining a memory buffer to allow training on independent / randomized minibatches of previous states. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1489f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent:\n",
    "    def __init__(self, window_size, is_eval=False, model_name=\"\"):\n",
    "        #State size depends and is equal to the the window size, n previous days\n",
    "        self.window_size = window_size\n",
    "        self.state_size = window_size # normalized previous days, \n",
    "        self.action_size = 3 # sit, buy, sell\n",
    "        self.memory = deque(maxlen=1000)\n",
    "        # inventory of close prices \n",
    "        self.inventory = []\n",
    "        self.model_name = model_name\n",
    "        self.is_eval = is_eval\n",
    "\n",
    "        self.gamma = 0.95\n",
    "        self.epsilon = 1.0\n",
    "        self.epsilon_min = 0.01\n",
    "        self.epsilon_decay = 0.995\n",
    "        \n",
    "        self.model = keras.models.load_model(model_name) if is_eval else self._model()\n",
    "\n",
    "    #Deep Q Learning model- returns the q-value when given state as input \n",
    "    def _model(self):\n",
    "        model = DQN(self.state_size, self.action_size).model\n",
    "        return model\n",
    "    \n",
    "    #Return the action on the value function\n",
    "    #With probability (1-$\\epsilon$) choose the action which has the highest Q-value.\n",
    "    #With probability ($\\epsilon$) choose any action at random.\n",
    "    #Intitially high epsilon-more random, later less\n",
    "    #The trained agents were evaluated by different initial random condition\n",
    "    #and an e-greedy policy with epsilon 0.05. This procedure is adopted to minimize the possibility of overfitting during evaluation.\n",
    " \n",
    "    def act(self, state): \n",
    "        #If it is test and self.epsilon is still very high, once the epsilon become low, there are no random\n",
    "        #actions suggested.\n",
    "        if not self.is_eval and random.random() <= self.epsilon:\n",
    "            return random.randrange(self.action_size)      \n",
    "        # print(\"state\", state)\n",
    "        options = self.model.predict(state.flatten().reshape(1, self.window_size))\n",
    "        #action is based on the action that has the highest value from the q-value function.\n",
    "        return np.argmax(options[0])\n",
    "\n",
    "    def expReplay(self, batch_size):\n",
    "        mini_batch = []\n",
    "        l = len(self.memory)\n",
    "        for i in range(l - batch_size + 1, l):\n",
    "            mini_batch.append(self.memory[i])\n",
    "        \n",
    "        # the memory during the training phase. \n",
    "        for state, action, reward, next_state, done in mini_batch:\n",
    "            target = reward # reward or Q at time t    \n",
    "            #update the Q table based on Q table equation\n",
    "            #set_trace()\n",
    "            if not done:\n",
    "                #max of the array of the predicted. \n",
    "                target = reward + self.gamma * np.amax(self.model.predict(next_state.flatten().reshape(1, self.window_size)))  \n",
    "                \n",
    "            # Q-value of the state currently from the table    \n",
    "            target_f = self.model.predict(state.flatten().reshape(1, self.window_size))  \n",
    "            # Update the output Q table for the given action in the table     \n",
    "            target_f[0][action] = target\n",
    "            #train and fit the model where state is X and target_f is Y, where the target is updated. \n",
    "            self.model.fit(state.flatten().reshape(1, self.window_size), target_f, epochs=1, verbose=0)\n",
    "        \n",
    "        if self.epsilon > self.epsilon_min:\n",
    "            self.epsilon *= self.epsilon_decay"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfee09a6",
   "metadata": {},
   "source": [
    "# 4. Train the Agent\n",
    "Now that our agent is defined, we are ready to train it. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a05cd0f4",
   "metadata": {},
   "source": [
    "### Helper Functions\n",
    "Before we define the training loop, we will write some helper functions: one for printing price data, one to define the sigmoind funtion, one to grab the state representation, and one to plot the output of our trained model. The printing, sigmoid, and plotting functions are defined for you. You must define the function which gets the state representation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "722fbbdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prints formatted price\n",
    "def formatPrice(n):\n",
    "    return (\"-$\" if n < 0 else \"$\") + \"{0:.2f}\".format(abs(n))\n",
    "\n",
    "# returns the sigmoid\n",
    "def sigmoid(x):\n",
    "    return 1 / (1 + math.exp(-x))\n",
    "\n",
    "# returns an an n-day state representation ending at time t\n",
    "\n",
    "def getState(data, t, n):    \n",
    "    d = t - n + 1\n",
    "    if d >= 0:\n",
    "        block = data[d:t + 1] \n",
    "    else:\n",
    "        block =  np.array([data[0]]*n) # pad with t0\n",
    "    res = []\n",
    "    for i in range(n - 1):\n",
    "        feature_res = []\n",
    "        for feature in range(data.shape[1]):\n",
    "            feature_res.append(sigmoid(block[i + 1, feature] - block[i, feature]))\n",
    "        res.append(feature_res)\n",
    "    # display(res)\n",
    "    return np.array([res])\n",
    "\n",
    "# Plots the behavior of the output\n",
    "def plot_behavior(data_input, states_buy, states_sell, profit, train=True):\n",
    "    fig = plt.figure(figsize = (15,5))\n",
    "    plt.plot(data_input, color='k', lw=2., label= 'Close Price')\n",
    "    plt.plot(data_input, '^', markersize=10, color='r', label = 'Buying signal', markevery = states_buy)\n",
    "    plt.plot(data_input, 'v', markersize=10, color='g', label = 'Selling signal', markevery = states_sell)\n",
    "    plt.title('Total gains: %f'%(profit))\n",
    "    plt.legend()\n",
    "    # locs, labels = plt.xticks()\n",
    "    # print(locs, labels)\n",
    "    if train:\n",
    "        plt.xticks(range(len(train_df.index.values)), train_df.index.values, rotation=45) # location, labels\n",
    "    else:\n",
    "        plt.xticks(range(len(test_df.index.values)), test_df.index.values, rotation=45) # location, labels\n",
    "\n",
    "    #plt.savefig('output/'+name+'.png')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e6bd8f6",
   "metadata": {},
   "source": [
    "### Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "798817f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# display the shape of your training data in order to remond yourself how may features and examples there are in your training set\n",
    "X_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0550bb8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "keras.utils.disable_interactive_logging()\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "window_size = 1\n",
    "agent = Agent(window_size)\n",
    "dot = keras.utils.model_to_dot(\n",
    "    agent.model,\n",
    "    show_shapes=True,\n",
    "    show_dtype=True,\n",
    "    show_layer_names=True,\n",
    ")\n",
    "dot.write(\"model.png\", format='png')\n",
    "from IPython import display\n",
    "\n",
    "display.Image('model.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "851650c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "keras.config.disable_traceback_filtering()\n",
    "\n",
    "# track number of examples in dataset (i.e. number of days to train on)\n",
    "l = X_train[:,0].shape[0] - 1\n",
    "\n",
    "# batch size defines how often to run the expReplay method\n",
    "batch_size = 32\n",
    "\n",
    "#An episode represents a complete pass over the data.\n",
    "episode_count = 2\n",
    "\n",
    "st = time.time()\n",
    "\n",
    "for e in range(episode_count + 1):\n",
    "    # print()\n",
    "    state = getState(X_train, 0, window_size + 1)\n",
    "    #set_trace()\n",
    "    total_profit = 0\n",
    "    agent.inventory = []\n",
    "    states_sell = []\n",
    "    states_buy = []\n",
    "    for t in tqdm(range(l), desc=\"Running episode \" + str(e) + \"/\" + str(episode_count)):\n",
    "        action = agent.act(state)    \n",
    "        # sit\n",
    "        next_state = getState(X_train, t + 1, window_size + 1)\n",
    "        reward = 0\n",
    "\n",
    "        if action == 1: # buy\n",
    "            # inverse transform to get true buy price in dollars\n",
    "            buy_price = X_train[t].item()\n",
    "            agent.inventory.append(buy_price)\n",
    "            # print('inventory', agent.inventory)\n",
    "            states_buy.append(t)\n",
    "            print(\"Buy: \" + formatPrice(buy_price))\n",
    "\n",
    "        elif action == 2 and len(agent.inventory) > 0: # sell\n",
    "            bought_price = agent.inventory.pop(0)  \n",
    "            # print('inventory', agent.inventory)\n",
    "            # inverse transform to get true sell price in dollars\n",
    "            sell_price = X_train[t].item()\n",
    "\n",
    "            # reward is max of profit (close price at time of sell - close price at time of buy)\n",
    "            reward = max(sell_price - bought_price, 0)\n",
    "            total_profit += sell_price - bought_price\n",
    "            states_sell.append(t)\n",
    "            print(\"Sell: \" + formatPrice(sell_price) + \" | Profit: \" + formatPrice(sell_price - bought_price))\n",
    "\n",
    "        done = True if t == l - 1 else False\n",
    "        #appends the details of the state action etc in the memory, which is used further by the exeReply function\n",
    "        agent.memory.append((state, action, reward, next_state, done))\n",
    "        state = next_state\n",
    "\n",
    "        if done:\n",
    "            print(\"--------------------------------\")\n",
    "            print(\"Total Profit: \" + formatPrice(total_profit))\n",
    "            print(\"--------------------------------\")\n",
    "            plot_behavior(X_train, states_buy, states_sell, total_profit)\n",
    "\n",
    "        if len(agent.memory) > batch_size:\n",
    "            agent.expReplay(batch_size)    \n",
    "            \n",
    "\n",
    "    if e % 2 == 0:\n",
    "        agent.model.save(\"model_ep\" + str(e) + \".keras\")\n",
    "\n",
    "print(\"TOTAL TRAINING TIME\", time.time()-st)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8adb831d",
   "metadata": {},
   "source": [
    "# 5. Test the trained model \n",
    "Finally, we get to test our trained model to see how well it performs in our test set. Using the training loop above, define a method to run our trained model on our X_test dataset. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cbd7885",
   "metadata": {},
   "source": [
    "## Define Parameters\n",
    "Some test parameters are defined for you below. Fill out the missing data. If you need a hint, look up at the training loop. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02b3f476",
   "metadata": {},
   "outputs": [],
   "source": [
    "l_test = len(X_test) - 1\n",
    "state = getState(X_test, 0, window_size + 1)\n",
    "total_profit = 0\n",
    "done = False\n",
    "states_sell_test = []\n",
    "states_buy_test = []\n",
    "#Get the trained model\n",
    "agent = Agent(window_size, is_eval=True, model_name=\"model_ep\"+str(episode_count)+\".keras\")\n",
    "agent.inventory = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b474ba5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for t in range(l_test):\n",
    "    action = agent.act(state)\n",
    "    #print(action)\n",
    "    #set_trace()\n",
    "    next_state = getState(X_test, t + 1, window_size + 1)\n",
    "    reward = 0\n",
    "\n",
    "    if action == 1: # buy\n",
    "        # inverse transform to get true buy price in dollars\n",
    "        buy_price = X_test[t].item()\n",
    "        agent.inventory.append(buy_price)\n",
    "        states_buy_test.append(t)\n",
    "        print(\"Buy: \" + formatPrice(buy_price))\n",
    "\n",
    "    elif action == 2 and len(agent.inventory) > 0: # sell\n",
    "        bought_price = agent.inventory.pop(0)  \n",
    "        # print('inventory', agent.inventory)\n",
    "        # inverse transform to get true sell price in dollars\n",
    "        sell_price = X_test[t].item()\n",
    "\n",
    "        # reward is max of profit (close price at time of sell - close price at time of buy)\n",
    "        reward = max(sell_price - bought_price, 0)\n",
    "        total_profit += sell_price - bought_price\n",
    "        states_sell_test.append(t)\n",
    "        print(\"Sell: \" + formatPrice(sell_price) + \" | Profit: \" + formatPrice(sell_price - bought_price))\n",
    "\n",
    "\n",
    "    if t == l_test - 1:\n",
    "        done = True\n",
    "        \n",
    "    agent.memory.append((state, action, reward, next_state, done))\n",
    "    state = next_state\n",
    "\n",
    "    if done:\n",
    "        print(\"------------------------------------------\")\n",
    "        print(\"Total Profit: \" + formatPrice(total_profit))\n",
    "        print(\"------------------------------------------\")\n",
    "        \n",
    "plot_behavior(X_test, states_buy_test, states_sell_test, total_profit, train=False)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rl-project-udacity",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
